fix: propagate re-registered node id and make rebind thread-safe by kaiitunnz · Pull Request #94 · mlsys-io/FlowMesh

kaiitunnz · 2026-07-05T14:07:07Z

Purpose

Follow-up to #91 / #93. #91 makes a node whose root-registry record is lost re-register under a fresh, server-allocated node id, and #93 rebinds the child-side subscriptions and re-homes workers. Two correctness gaps remained: the rebind mutated a live redis-py PubSub from the heartbeat thread while the reader thread was inside listen() (unsafe cross-thread access, made reachable by #92's health-check PINGs on that connection), and the parent server process kept stale copies of the node id after re-registration. This PR closes both while keeping the id server-allocated (no reuse).

Changes

src/server/supervisor/services/task_listener.py, command_listener.py — rebind now records the target node id under a lock; the pubsub reader thread applies the subscribe/unsubscribe between get_message polls and signals completion via an event, so no PubSub is ever mutated off the reader thread.
src/server/clients/redis.py — extract parse_pubsub_message so the blocking and polling readers share one frame parser, and add an eval wrapper (sync + async) for the re-home script.
src/server/supervisor/services/grpc_server.py — guard rebind_node and the RegisterWorker node-id/registry critical section with a lock, and re-home workers with a single conditional-HSET Lua EVAL that skips (rather than resurrects) workers whose record was evicted.
src/server/supervisor/supervisor.py, src/server/main.py — the child re-puts the new id on the handshake queue (non-blocking) and a parent watcher thread refreshes app.state.node_id and the EventMonitor's own-node through listeners.
tests/server/ — reader-owned rebind, worker re-home, parse_pubsub_message, and parent node-id watcher coverage, with shared re-register stubs in supervisor_helpers.py.

Design

The id stays server-allocated and fresh on each re-register — reusing the old id was rejected because a wiped root resets NODE_ID_SEQ (collision risk) and accepting a client-chosen id weakens the allocation invariant. The new id is instead propagated coherently: subscriptions move on the thread that owns the socket, and the parent learns the id over the existing multiprocessing handshake queue rather than the telemetry stream (the queue is in-process and reliable; telemetry is exactly what intermediaries cull, per #92). _on_reregister rebinds the subscriptions and waits for the reader to apply them before re-homing workers, so the dispatcher never publishes to a channel nobody is listening on yet.

Test Plan

uv run pytest tests/ --ignore=tests/worker/test_mp_executor_cleanup_gpu.py
uv run pre-commit run --all-files

End-to-end on a local root + a CPU worker on a freshly built image carrying this branch: delete the node's control-plane keys to simulate the root registry losing it, then confirm the node re-registers under a new id within one heartbeat interval and stays functional — the new id answers node commands and a workflow submitted afterward runs to DONE (dispatch rebind + worker re-home).

Test Result

$ uv run pytest tests/ --ignore=tests/worker/test_mp_executor_cleanup_gpu.py   # 1169 passed
$ uv run pre-commit run --all-files   # gitleaks / isort / black / ruff / codespell / mypy / sync-requirements passed

End-to-end: after the node's keys were deleted it re-registered nde-1 → nde-2 within one heartbeat interval; node worker list <new_id> answered on the new command channel (no hang), and a workflow submitted after re-register ran to DONE with the worker observed homed under the new id.

Pre-submission Checklist

I have read the contribution guidelines.
I have run pre-commit run --all-files and fixed any issues.
I have added or updated tests covering my changes.
I have verified that uv run pytest tests/ passes locally.
If I changed shared schemas or proto definitions, I have checked downstream compatibility across Server and Worker.
If I changed the SDK or CLI, I have verified the affected packages work.
If this is a breaking change, I have prefixed the PR title with [BREAKING] and described migration steps above.
I have updated documentation or config examples if user-facing behavior changed.

Split the decode/type-filter body out of iter_pubsub_messages so a polling reader (get_message) and the blocking reader (listen) share one frame parser. Behavior of iter_pubsub_messages is unchanged. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

The dispatch and command subscriptions were re-subscribed directly from the heartbeat thread while the reader thread was inside pubsub.listen(). A redis-py PubSub is not safe to mutate from a second thread, and the health-check PINGs on that connection make interleaved socket writes reachable. rebind() now only records the target id under a lock; the reader polls with get_message and applies the subscribe/unsubscribe itself. wait_rebound lets the caller order work after the switch. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

rebind_node ran on the heartbeat thread and mutated _node_id plus the in-process registry while RegisterWorker ran on the grpc loop, unlocked; a worker registering mid-rebind could be homed on the stale node. Guard _node_id and the register/re-home window with a lock, skip workers whose Redis record was evicted rather than resurrecting a partial record, and pipeline the rewrites so the lock spans two round trips, not N. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

A re-register mints a fresh node id inside the child, but the parent only read it once via the startup handshake, leaving app.state.node_id (request auth scope) and the EventMonitor's own-node (shutdown correlation) stale. The child now re-puts the new id on the handshake queue and a parent watcher thread refreshes the cached id and notifies listeners. The put is non-blocking so it cannot stall the heartbeat thread. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

Update the re-register tests to the pending-rebind contract (rebind records the target and touches no pubsub; the reader-side apply moves the subscription) and add the cross-thread regression guard, the worker re-home exists guard, parse_pubsub_message decoding, and the parent node-id watcher/listener path. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

The two-pipeline existence-check-then-HSET could recreate a partial worker record if the key was deleted between the checks. Do the conditional rewrite in a single Lua EVAL so a worker deleted mid-rebind is skipped, not resurrected. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

Register the re-register callback once the dispatch/command reader threads are running, so a re-register cannot block on wait_rebound waiting for a reader that has not started yet. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

…client Tighten the docstrings/comments added by the re-register work, group the node-id watcher methods under their own section, and add the async twin of SyncRedisClient.eval to keep the two clients at parity. No behavior change. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

Extract the StubRegistry/StubLifecycle test doubles into a shared supervisor_helpers module, consolidate the duplicated _reregister_if_lost tests into the lifecycle suite, and build the fakes via real constructors / subclasses so no test needs a type: ignore. Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

timzsu

Two correctness-related comments, plus several comments on the clean-up of pre-existing typing issues.

timzsu · 2026-07-07T07:05:33Z

+    # These constants are Linux-only; skip any the running platform lacks (e.g.
+    # macOS/Windows dev hosts) so importing this module there doesn't fail.


Do we plan to support macOS/Windows? I thought Linux is our only target platform

timzsu · 2026-07-07T08:48:31Z

+        task_listener.wait_rebound(_REBIND_APPLY_TIMEOUT_SEC)
+        command_listener.wait_rebound(_REBIND_APPLY_TIMEOUT_SEC)


If the wait times out, then the server is rebound to the new ID without a proper Redis channel configured. I think we should capture the results of the wait, and if they are False, we should retry properly.

timzsu · 2026-07-07T08:59:13Z

+        # the heartbeat thread, which must not stall on a full queue).
+        try:
+            node_id_queue.put_nowait(new_node_id)
+        except QueueFull:


How about draining the oldest item from the queue if the queue is full?

timzsu · 2026-07-07T09:17:01Z

            return value[7:]  # type: ignore
        if key_l == "x-worker-token":
            return value  # type: ignore


I think we should remove these two type: ignores, which were from the first commit.

timzsu · 2026-07-07T09:17:57Z

    def _get_worker_from_context(
        self, context: grpc.aio.ServicerContext
    ) -> WorkerAdapter | None:
        token = _token_from_metadata(context.invocation_metadata())  # type: ignore


Similarly, I think this type: ignore can also be fixed easily.

timzsu · 2026-07-07T09:18:09Z

    def _get_worker_id_from_context(
        self, context: grpc.aio.ServicerContext
    ) -> str | None:
        token = _token_from_metadata(context.invocation_metadata())  # type: ignore


Same as above.

timzsu · 2026-07-07T09:19:17Z



 def _sync[T](value: Awaitable[T] | T) -> T:
    return value  # type: ignore


Can we remove this # type: ignore?

timzsu · 2026-07-07T09:21:19Z

    try:
        rm = ResourceManager.get_instance()
        max_gpu_count = rm.total_gpu_count
        current_gpu_count_getter = lambda: rm.available_gpu_count  # noqa: E731


I think this noqa is also not justified.

kaiitunnz added 9 commits July 5, 2026 21:29

kaiitunnz marked this pull request as ready for review July 6, 2026 15:07

kaiitunnz requested a review from timzsu July 6, 2026 15:07

timzsu requested changes Jul 7, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix: propagate re-registered node id and make rebind thread-safe#94

fix: propagate re-registered node id and make rebind thread-safe#94
kaiitunnz wants to merge 9 commits into
mainfrom
fix/node-reregister-propagation

kaiitunnz commented Jul 5, 2026 •

edited

Loading

Uh oh!

timzsu left a comment

Uh oh!

timzsu Jul 7, 2026

Uh oh!

timzsu Jul 7, 2026

Uh oh!

timzsu Jul 7, 2026

Uh oh!

timzsu Jul 7, 2026

Uh oh!

timzsu Jul 7, 2026

Uh oh!

timzsu Jul 7, 2026

Uh oh!

timzsu Jul 7, 2026

Uh oh!

timzsu Jul 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		# These constants are Linux-only; skip any the running platform lacks (e.g.
		# macOS/Windows dev hosts) so importing this module there doesn't fail.

		task_listener.wait_rebound(_REBIND_APPLY_TIMEOUT_SEC)
		command_listener.wait_rebound(_REBIND_APPLY_TIMEOUT_SEC)



		def _sync[T](value: Awaitable[T] \| T) -> T:
		return value # type: ignore

Uh oh!

Conversation

kaiitunnz commented Jul 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Changes

Design

Test Plan

Test Result

Uh oh!

timzsu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kaiitunnz commented Jul 5, 2026 •

edited

Loading